package com.spark.sniffer.tracing;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * 线程上下文
 * 
 * @author Cobe
 *
 */
public class TracingContext {
    private static ThreadLocal<TracingContext> CONTEXT = new ThreadLocal<TracingContext>() {
        @Override
        protected TracingContext initialValue() {
            return new TracingContext();
        }
    };

    private String traceId = null;

    // 存储透传数据
    private Map<String, String> passThroughMap = null;

    // 存储其他临时数据
    private Map<String, Object> dataHolder = null;

    private Span entranceSpan = null;
    private LinkedList<Span> activeSpanStack;
    private List<Span> finishedSpanStack;

    public static TracingContext getContext() {
        return CONTEXT.get();
    }

    public static void removeContext() {
        CONTEXT.remove();
    }

    public String getTraceId() {
        return traceId;
    }

    public TracingContext setTraceId(String traceId) {
        this.traceId = traceId;
        return this;
    }

    public TracingContext putPassThrough(String key, String value) {
        if (value != null && value.length() > 0) {
            if (passThroughMap == null) {
                passThroughMap = new HashMap<>();
            }
            passThroughMap.put(key, value);
        }
        return this;
    }

    public Map<String, String> getPassThroughMap() {
        return passThroughMap;
    }

    public TracingContext entrance(Span.Type type, String name) {
        entranceSpan = new Span(type, name);
        activeSpanStack = new LinkedList<Span>();
        finishedSpanStack = new LinkedList<Span>();
        return this;
    }

    public Span stopEntrance(int code) {
        // entranceSpan.getEndTime()>0是因为在handleMethodException中已设置，因此在afterMethod不再设置
        if (entranceSpan != null && entranceSpan.getEndTime() == 0) {
            entranceSpan.setEndTime(System.currentTimeMillis()).setCode(code);
        }
        return entranceSpan;
    }

    public Span getEntrance() {
        return entranceSpan;
    }

    public TracingContext startSpan(Span.Type type, String name) {
        // 长度限制防止出现意外循环导致占用内存过大
        // 不过长度超过限制时，stopSpan出栈时将会出现混乱
        if (activeSpanStack != null && activeSpanStack.size() < 500) {
            activeSpanStack.push(new Span(type, name));
        }
        return this;
    }

    public Span stopSpan() {
        if (activeSpanStack == null || activeSpanStack.isEmpty()) {
            return null;
        }

        Span span = activeSpanStack.pop().setEndTime(System.currentTimeMillis());
        finishedSpanStack.add(span);
        return span;
    }

    public Span setActiveSpanCode(int code) {
        return activeSpanStack == null || activeSpanStack.isEmpty() ? null : activeSpanStack.peek().setCode(code);
    }

    public Map<String, Object> getDataHolder() {
        return dataHolder;
    }

    public <T> T getData(String key) {
        return dataHolder == null ? null : (T)dataHolder.get(key);
    }

    public <T> T getData(String key, T defaultValue) {
        return dataHolder == null ? defaultValue : (T)dataHolder.getOrDefault(key, defaultValue);
    }

    public TracingContext putData(String key, Object value) {
        if (value == null) {
            return this;
        }

        if (this.dataHolder == null) {
            this.dataHolder = new HashMap<>();
        }

        this.dataHolder.put(key, value);

        return this;
    }

    public TracingContext removeData(String key) {
        if (this.dataHolder == null) {
            return this;
        }

        this.dataHolder.remove(key);

        return this;
    }

    /**
     * 调用栈转成字符串，格式如：a->b->
     * 
     * @return
     */
    public String callstack2string() {
        if (finishedSpanStack == null || finishedSpanStack.isEmpty()) {
            return "";
        }

        long entranceTime = this.getEntrance().getBeginTime();

        StringBuilder buffer = new StringBuilder();
        for (int i = 0; i < finishedSpanStack.size(); i++) {
            Span span = finishedSpanStack.get(i);
            buffer.append(span.getType().value());
            if (span.getName() != null) {
                buffer.append("(").append(span.getName()).append(")");
            }
            buffer.append(":").append(span.getDuration()).append(":").append(span.getEndTime() - entranceTime);
            if (!span.isSuccess()) {
                buffer.append(":x");
            }
            buffer.append("->");
        }
        buffer.delete(buffer.length() - 2, buffer.length());

        return buffer.toString();
    }

    /**
     * 透传数据转成字符串，格式如：abtest_group:grey|X-User-Agent:Android 1.0.0
     * 
     * @return
     */
    public String passThrough2string() {
        // 为空或只有traceid不打印
        if (passThroughMap == null || passThroughMap.size() == 1) {
            return "";
        }

        StringBuilder buffer = new StringBuilder();
        for (Map.Entry<String, String> entry : passThroughMap.entrySet()) {
            if (TraceId.KEY.equals(entry.getKey())) {
                continue;
            }
            buffer.append(entry.getKey()).append(":").append(entry.getValue()).append("|");
        }
        buffer.deleteCharAt(buffer.length() - 1);
        return buffer.toString();
    }
}
